package leetcode

type TreeNode struct {
	Val   int
	Left  *TreeNode
	Right *TreeNode
}

func levelOrder1(root *TreeNode) [][]int {
	return traverse(root, 1, [][]int{})
}

func traverse(root *TreeNode, level int, rs [][]int) [][]int {
	if root == nil {
		return rs
	}

	if level > len(rs) {
		rs = append(rs, []int{})
	}
	rs[level-1] = append(rs[level-1], root.Val)

	rs = traverse(root.Left, level+1, rs)
	rs = traverse(root.Right, level+1, rs)
	return rs
}

func levelOrder(root *TreeNode) [][]int {
	if root == nil {
		return [][]int{}
	}

	if root.Left == nil && root.Right == nil {
		return [][]int{[]int{root.Val}}
	}

	rs, level, n := [][]int{}, 0, new(TreeNode)
	q, next := []*TreeNode{root}, make([]*TreeNode, 0, 2)
	i := 0
	for {
		n, i = q[i], i+1
		if len(rs) < level+1 {
			rs = append(rs, []int{n.Val})
		} else {
			rs[level] = append(rs[level], n.Val)
		}

		if n.Left != nil {
			next = append(next, n.Left)
		}
		if n.Right != nil {
			next = append(next, n.Right)
		}

		if i < len(q) {
			continue
		}

		if len(next) == 0 {
			break
		}

		q, next, level, i = next, q[:0], level+1, 0
	}

	return rs
}
